/**
 * ============================================================================
 *
 * Copyright (C) 2018, Hisilicon Technologies Co., Ltd. All Rights Reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *   1 Redistributions of source code must retain the above copyright notice,
 *     this list of conditions and the following disclaimer.
 *
 *   2 Redistributions in binary form must reproduce the above copyright notice,
 *     this list of conditions and the following disclaimer in the documentation
 *     and/or other materials provided with the distribution.
 *
 *   3 Neither the names of the copyright holders nor the names of the
 *   contributors may be used to endorse or promote products derived from this
 *   software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 * ============================================================================
 */
#include <vector>
#include <sstream>
#include <cmath>
#include <sys/time.h>
#include <opencv2/opencv.hpp>
#include "hiaiengine/log.h"
#include "format_convertion.h"
#include "util.hpp"

using hiai::Engine;
using namespace std::__cxx11;

// register data type
HIAI_REGISTER_DATA_TYPE("EngineTransT", EngineTransT);
HIAI_REGISTER_DATA_TYPE("OutputT", OutputT);
HIAI_REGISTER_DATA_TYPE("NewImageParaT", NewImageParaT);

// constants
namespace {
// output port (engine port begin with 0)
const uint32_t kSendDataPort = 0;

// style transfer model output width
const uint32_t kModelOutputImageWidth = 540;

// style transfer model output height
const uint32_t kModelOutputImageHeight = 360;
}

FormatConvertion::FormatConvertion() {
  format_config_ = nullptr;
}

HIAI_StatusT FormatConvertion::Init(
    const hiai::AIConfig& config,
    const std::vector<hiai::AIModelDescription>& model_desc) {
  HIAI_ENGINE_LOG("Begin initialize!");

  // get configurations
  if (format_config_ == nullptr) {
    format_config_ = std::make_shared<ResizedImageConfig>();
  }

  // get parameters from graph.config
  for (int index = 0; index < config.items_size(); index++) {
    const ::hiai::AIConfigItem& item = config.items(index);
    const std::string& name = item.name();
    const std::string& value = item.value();

    if (name == "resized_width") {
      format_config_->resized_width = atoi(value.data());
    } else if (name == "resized_height") {
      format_config_->resized_height = atoi(value.data());
    } else {
      HIAI_ENGINE_LOG("unused config name: %s", name.c_str());
    }
  }

  HIAI_ENGINE_LOG(HIAI_DEBUG_INFO, "End initialize!");
  return HIAI_OK;
}

HIAI_StatusT FormatConvertion::SendResults(
    const std::shared_ptr<EngineTransT> &inference_res) {
  HIAI_StatusT status = HIAI_OK;
  std::shared_ptr<EngineTransT> trans_data = std::make_shared<EngineTransT>();
  trans_data->b_info = inference_res->b_info;
  trans_data->status = true;

  std::vector<OutputT> output_data_vec = inference_res->output_datas;
  // dealing every image
  std::vector<NewImageParaT> converted_jpeg;
  for (uint32_t ind = 0; ind < inference_res->b_info.batch_size; ind++) {
    OutputT out = output_data_vec[ind];
    int out_size = out.size;
    std::shared_ptr<u_int8_t> out_data(new u_int8_t[out_size]);
    memcpy_s(out_data.get(), out_size, out.data.get(), out_size);

    //**** Firstly, we split chw format output RGB channel alone, then we merge then together to form hwc bgr image.
    unsigned long hwc_resize_start_time = get_current_time();
    cv::Mat f32_mat[3];
    f32_mat[0].create(kModelOutputImageHeight, kModelOutputImageWidth, CV_32FC1);
    f32_mat[1].create(kModelOutputImageHeight, kModelOutputImageWidth, CV_32FC1);
    f32_mat[2].create(kModelOutputImageHeight, kModelOutputImageWidth, CV_32FC1);
    f32_mat[0].data = out_data.get();
    f32_mat[1].data = out_data.get() + out.size/3;
    f32_mat[2].data = out_data.get() + 2*out.size/3;
    cv::Mat u8_mat[3];
    f32_mat[0].convertTo(u8_mat[0], CV_8UC1);
    f32_mat[1].convertTo(u8_mat[1], CV_8UC1);
    f32_mat[2].convertTo(u8_mat[2], CV_8UC1);

    // Merge the three channel again to form bgr hwc image, and resize them to 1080*720
    vector<cv::Mat> channels;
    channels.push_back(u8_mat[2]);
    channels.push_back(u8_mat[1]);
    channels.push_back(u8_mat[0]);
    cv::Mat src_mat, dst_mat;
    merge(channels, src_mat);
    resize(src_mat, dst_mat, cv::Size(format_config_->resized_width, format_config_->resized_height));
    unsigned long hwc_resize_time = get_current_time() - hwc_resize_start_time;

    // Here, convert BGR to I420
    unsigned long brg2I420_start_time = get_current_time();
    cv::Mat image_yuv;
    cvtColor(dst_mat, image_yuv, CV_BGR2YUV_I420);
    int i420_image_size = format_config_->resized_width * format_config_->resized_height * 3 / 2;
    std::shared_ptr<u_int8_t> i420_image_data(new u_int8_t[i420_image_size]);
    memcpy_s(i420_image_data.get(), i420_image_size, image_yuv.data, i420_image_size);
    unsigned long brg2I420_time = get_current_time() - brg2I420_start_time;
    HIAI_ENGINE_LOG(HIAI_ENGINE_RUN_ARGS_NOT_RIGHT, "format convert.resize time:%d (us) bgr2I420 time:%d (us)", hwc_resize_time, brg2I420_time);

    OutputT out_i420;
    out_i420.size = i420_image_size;
    out_i420.data = i420_image_data;
    trans_data->output_datas.push_back(out_i420);
  }
  // send results and original image data to post process (port 0)
  HIAI_StatusT hiai_ret = SendData(kSendDataPort, "EngineTransT",
                                   std::static_pointer_cast<void>(trans_data));
  
  return hiai_ret;
}

HIAI_IMPL_ENGINE_PROCESS("format_convertion",
    FormatConvertion, INPUT_SIZE) {
  // check arg0 is null or not
  if (arg0 == nullptr) {
    HIAI_ENGINE_LOG(HIAI_ENGINE_RUN_ARGS_NOT_RIGHT,
                    "Failed to process invalid message.");
    return HIAI_ERROR;
  }

  // check original image is empty or not
  std::shared_ptr<EngineTransT> inference_res = std::static_pointer_cast<
      EngineTransT>(arg0);

  // dealing inference results
  return SendResults(inference_res);
}
